

import sys
sys.path.insert(0, ".")
sys.path.insert(0, "..")

import czsc
print(czsc.__version__)

from datetime import datetime
from typing import List
import traceback
from tqdm import tqdm
import pandas as pd
import numpy as np
from tqdm import tqdm_notebook as tqdm
from czsc.analyze import KlineAnalyze
from schedule.utils.query_data import get_data

def selector_bi01(kline):
    """笔中枢三买选股"""
    res = {"is_match": False, "notes": None, "dt": None, "price": None,
       "name": "selector_bi01", "symbol": None, "desc": "笔中枢三买选股"}
    try:
        ka = KlineAnalyze(kline, ma_params=(5, 34, 120, 233), bi_mode="new")
        points = ka.bi_list[-6:]

        if len(points) == 6 and points[-1]['fx_mark'] == "d":
            zs_g = min([x['bi'] for x in points[:4] if x['fx_mark'] == 'g'])
            zs_d = max([x['bi'] for x in points[:4] if x['fx_mark'] == 'd'])

            if points[-1]['bi'] > zs_g > zs_d:
                res['symbol'] = ka.symbol
                res['is_match'] = True
                res['dt'] = ka.end_dt
                res['price'] = ka.latest_price

    except:
        print("{} 分析失败".format(symbol))
        traceback.print_exc()
    return res

def selector_bi02(kline):
    """三笔回调构成三买"""
    res = {"is_match": False, "notes": None, "dt": None, "price": None,
           "name": "selector_bi02", "symbol": None, "desc": "三笔回调构成三买"}

    try:
        ka = KlineAnalyze(kline, ma_params=(5, ), bi_mode="new")
        if len(ka.bi_list) > 16 and ka.bi_list[-1]['fx_mark'] == 'd':
            points = ka.bi_list[-16:]
            fd1_points = [x['bi'] for x in points[0: 4]]
            fd2_points = [x['bi'] for x in points[3: 7]]
            fd3_points = [x['bi'] for x in points[6: 10]]
            fd4_points = [x['bi'] for x in points[9: 13]]
            fd5_points = [x['bi'] for x in points[12: 16]]
            fd1 = {"high": max(fd1_points), "low": min(fd1_points)}
            fd2 = {"high": max(fd2_points), "low": min(fd2_points)}
            fd3 = {"high": max(fd3_points), "low": min(fd3_points)}
            fd4 = {"high": max(fd4_points), "low": min(fd4_points)}
            fd5 = {"high": max(fd5_points), "low": min(fd5_points)}

            zs_g = min(fd1['high'], fd2['high'], fd3['high'])
            zs_d = max(fd1['low'], fd2['low'], fd3['low'])
            if zs_d < zs_g < fd5['low'] < fd4['high'] and fd4['high'] > max(fd1['high'], fd3['high']):
                res['symbol'] = ka.symbol
                res['is_match'] = True
                res['dt'] = ka.end_dt
                res['price'] = ka.latest_price
    except:
        traceback.print_exc()
    return res

def get_index_stocks():
    db = 'stock'
    sql = "select timestamp,codes,strategy,start,end,reason from stocks where strategy = 'good_company_roe_profit'  "  # and start<=%s and end>=%s
    arg = []
    stock_data = get_data(db=db, sql=sql, arg=arg)
    return stock_data

def get_klines(codes,end_date,freq,count):
    pass

# 使用聚宽数据在创业板综指上选股

symbols = get_index_stocks()
freq = "30min"

results = []
for symbol in tqdm(symbols, desc="selector"):
    try:
        kline = get_kline(symbol, end_date=datetime.now(), freq=freq, count=2000)
        for selector in [selector_bi01, selector_bi02]:
            try:
                res = selector(kline)
                if res['is_match']:
                    results.append(res)
            except:
                continue
    except:
        traceback.print_exc()
        print("fail on: {}".format(symbol))

print("选股结果：")
df = pd.DataFrame(results)
print(df)
# df.to_excel("选股结果.xlsx", index=False)